from exceptiongroup import catch
from langchain_core.messages import HumanMessage
from param.ipython import message

from langgraph_coder.src.agent.model.graph_state import GraphState


class CheckNode:
    """
    检查代码结点
    """
    def __call__(self, state: GraphState):
        """
        检查代码是否正确
        :param state: 图结构
        :return:
        """
        print("---代码检查结点---")
        messages = state["messages"]
        code_solution = state["generation"]

        # 拿到代码部分
        imports = code_solution["imports"]
        code = code_solution["code"]

        # 判断代码部分是否能够正确执行，先检查导入，后检查代码
        # 检查导入部分
        try:
            exec(imports)
        except Exception as e:
            print("---生成代码的Import部分失败---")
            error_message = HumanMessage(content=f"您的解决方案未能通过导入测试: {e}")
            messages += [error_message]
            # 更新图状态
            state["error"] = "yes"
            state["messages"] = messages
            return state

        # 检查代码部分
        try:
            exec(code)
        except Exception as e:
            print("---生成代码的Code Block部分失败---")
            error_message = HumanMessage(content=f"您的解决方案未能通过代码测试: {e}")
            messages += [error_message]
            # 更新图状态
            state["error"] = "yes"
            state["messages"] = messages
            return state

        print("---生成代码成功---")
        state["error"] = "no"
        return state
